from operator import add
from typing import TypedDict, Annotated

from langchain_core.messages import AnyMessage, HumanMessage
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.config import get_stream_writer
from langgraph.constants import START, END
from langgraph.graph import StateGraph

nodes = ["Supervisor", "Imperativequery", "ExploratoryAnalysis", "Dataextraction", "other"]

from zai import ZhipuAiClient
client = ZhipuAiClient(api_key="db2e03cca9654d4dba9b91b123548374.TftHaYZq4ydQPDjP")

class State(TypedDict):
    messages: Annotated[list[AnyMessage], add]
    type: str

def other_node(state: State):
    print("other node")
    writer = get_stream_writer()
    writer({"node", "other_node"})

    return {"messages": [HumanMessage(content="我暂时无法回答这个问题")], "type": "other"}

def supervisor_node(state: State):
    print("Supervisor node")
    writer = get_stream_writer()
    writer({"node", "Supervisor_node"})
    #根据用户的问题，对问题进行分类。分类结果保存到type当中
    prompt = """你是一个专业的SQL分类助手，负责对用户的问题关键词进行分类，并将任务分给其他的Agent执行。
可选的分类包括：
- Imperativequery：指令式查询，包括事实查询(获取具体数据)、状态查询(当前状态)、简单统计(计数、求和等)
- ExploratoryAnalysis：探索式分析，包括数据探索(浏览数据)、模式发现(趋势分析)、假设验证(测试假设)
- Dataextraction：数据提取，包括记录查询(获取记录列表)、数据导出(批量提取)、报表生成(格式化报表)
- other：其他问题

请根据用户输入的关键词进行分类，只返回分类名称（Imperativequery、ExploratoryAnalysis、Dataextraction、other），不要返回其他任何内容。不要解释。
用户问题："""
    
    # 获取用户消息内容
    user_message = state["messages"][-1].content if state["messages"] and hasattr(state["messages"][-1], 'content') else str(state["messages"][-1])
    
    full_prompt = prompt + user_message
    
    prompts = [
        {"role": "user", "content": full_prompt}
    ]
    
    # 如果已经有type属性了，表示问题已经交由其他节点处理完成了，就可以直接返回
    if "type" in state and state["type"] != "Supervisor":
        writer({"supervisor_step", f"已获得{state['type']}智能体处理结果"})
        return {"type": END}
    else:
        response = client.chat.completions.create(
            model="glm-4.5v",  # 使用更稳定的模型版本
            messages=prompts,
            temperature=0
        )
        # 正确访问响应内容
        typeRes = response.choices[0].message.content.strip()
        print(f"AI返回的分类结果: '{typeRes}'")  # 调试输出
        
        writer({"supervisor_step", f"问题分类结果：{typeRes}"})
        
        # 更灵活的分类匹配
        typeRes_lower = typeRes.lower()
        if "imperativequery" in typeRes_lower or "指令" in typeRes_lower:
            return {"type": "Imperativequery"}
        elif "exploratoryanalysis" in typeRes_lower or "探索" in typeRes_lower:
            return {"type": "ExploratoryAnalysis"}
        elif "dataextraction" in typeRes_lower or "数据提取" in typeRes_lower or "提取" in typeRes_lower:
            return {"type": "Dataextraction"}
        elif "other" in typeRes_lower or "其他" in typeRes_lower:
            return {"type": "other"}
        else:
            # 如果无法识别，默认返回other
            print(f"无法识别的分类: {typeRes}，默认返回other")
            return {"type": "other"}
    return {}

def Imperativequery_node(state: State):
    print("Imperativequery node")
    writer = get_stream_writer()
    writer({"node": "Imperativequery_node"})

    return {"messages": [HumanMessage(content="这是指令式查询节点处理结果")], "type": "Imperativequery"}

def ExploratoryAnalysis_node(state: State):
    print("ExploratoryAnalysis node")
    writer = get_stream_writer()
    writer({"node": "ExploratoryAnalysis_node"})

    return {"messages": [HumanMessage(content="这是探索式分析节点处理结果")], "type": "ExploratoryAnalysis"}

def Dataextraction_node(state: State):
    print("Dataextraction node")
    writer = get_stream_writer()
    writer({"node": "Dataextraction_node"})

    return {"messages": [HumanMessage(content="这是数据提取节点处理结果")], "type": "Dataextraction"}

#条件路由
def routing_func(state: State):
    if state["type"] == "Dataextraction":
        return "Dataextraction_node"
    elif state["type"] == "ExploratoryAnalysis":
        return "ExploratoryAnalysis_node"
    elif state["type"] == "Imperativequery":
        return "Imperativequery_node"
    elif state["type"] == END:
        return END
    else:
        return "other_node"


#构建图
builder = StateGraph(State)
#添加节点
builder.add_node("Supervisor_node", supervisor_node)
builder.add_node("Imperativequery_node", Imperativequery_node)
builder.add_node("ExploratoryAnalysis_node", ExploratoryAnalysis_node)
builder.add_node("Dataextraction_node", Dataextraction_node)
builder.add_node("other_node", other_node)
#添加边（edge）
builder.add_edge(START, "Supervisor_node")
builder.add_conditional_edges("Supervisor_node", routing_func, ["Imperativequery_node","ExploratoryAnalysis_node","Dataextraction_node","other_node",END])
builder.add_edge("Imperativequery_node", "Supervisor_node")
builder.add_edge("ExploratoryAnalysis_node", "Supervisor_node")
builder.add_edge("Dataextraction_node", "Supervisor_node")
builder.add_edge("other_node", "Supervisor_node")

#创建流程图graph
checkpointer = InMemorySaver()
graph = builder.compile(checkpointer=checkpointer)

#执行任务测试代码
if __name__ == "__main__":
    config = {
        "configurable":{
            "thread_id": "1"
        }
    }

    # 先测试一个简单的问题
    test_questions = [
         "查询今天的销售额",  # 应该分类为 Imperativequery
        # "分析销售趋势",      # 应该分类为 ExploratoryAnalysis  
        # "导出用户数据",      # 应该分类为 Dataextraction
        # "你好"              # 应该分类为 other
    ]
    
    for question in test_questions:
        print(f"\n=== 测试问题: {question} ===")
        try:
            for chunk in graph.stream(
                {"messages": [HumanMessage(content=question)]},
                config=config,
                stream_mode="custom"
            ):
                print(chunk)
        except Exception as e:
            print(f"错误: {e}")